#大数据2201刘源媛
import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import talib as ta
from datetime import datetime, timedelta
import warnings
import matplotlib.font_manager as fm

warnings.filterwarnings('ignore')

# 设置中文字体 - 更全面的字体设置
plt.rcParams['font.family'] = ['Arial Unicode MS', 'Microsoft YaHei', 'SimHei', 'WenQuanYi Micro Hei']  # 添加更多备选字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 设置图表样式
plt.style.use('seaborn-whitegrid')
sns.set(font='Microsoft YaHei')  # 设置seaborn的字体

# 自定义PO函数示例
def PO(close):
    """价格振荡指标"""
    return close - close.shift(1)

# K线形态识别函数
def identify_candlestick_patterns(df):
    """识别常见的K线形态"""
    result = df.copy()

    # 早晨之星
    result['MORNINGSTAR'] = ta.CDLMORNINGSTAR(result['open'], result['high'], result['low'], result['close'])

    # 黄昏之星
    result['EVENINGSTAR'] = ta.CDLEVENINGSTAR(result['open'], result['high'], result['low'], result['close'])

    # 十字星
    result['DOJI'] = ta.CDLDOJI(result['open'], result['high'], result['low'], result['close'])

    # 锤子线
    result['HAMMER'] = ta.CDLHAMMER(result['open'], result['high'], result['low'], result['close'])

    # 上吊线
    result['HANGINGMAN'] = ta.CDLHANGINGMAN(result['open'], result['high'], result['low'], result['close'])

    # 射击之星
    result['SHOOTINGSTAR'] = ta.CDLSHOOTINGSTAR(result['open'], result['high'], result['low'], result['close'])

    # 乌云盖顶
    result['DARKCLOUDCOVER'] = ta.CDLDARKCLOUDCOVER(result['open'], result['high'], result['low'], result['close'])

    # 看涨吞没
    result['ENGULFING_BULL'] = ta.CDLENGULFING(result['open'], result['high'], result['low'], result['close'])

    # 看跌吞没
    result['ENGULFING_BEAR'] = -ta.CDLENGULFING(result['open'], result['high'], result['low'], result['close'])

    # 计算每种形态的出现次数
    pattern_counts = {
        'Morning Star': (result['MORNINGSTAR'] != 0).sum(),
        'Evening Star': (result['EVENINGSTAR'] != 0).sum(),
        'Doji': (result['DOJI'] != 0).sum(),
        'Hammer': (result['HAMMER'] != 0).sum(),
        'Hanging Man': (result['HANGINGMAN'] != 0).sum(),
        'Shooting Star': (result['SHOOTINGSTAR'] != 0).sum(),
        'Dark Cloud': (result['DARKCLOUDCOVER'] != 0).sum(),
        'Bull Engulf': (result['ENGULFING_BULL'] > 0).sum(),
        'Bear Engulf': (result['ENGULFING_BEAR'] > 0).sum()
    }

    return result, pattern_counts

# 获取股票列表
def get_stock_list():
    pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
    df1 = pro.stock_basic(**{
        "ts_code": "",
        "name": "",
        "exchange": "",
        "market": "",
        "is_hs": "",
        "list_status": "",
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",
        "symbol",
        "name",
        "area",
        "industry",
        "cnspell",
        "market",
        "list_date",
        "act_name",
        "act_ent_type"
    ])
    return df1

# 获取日线行情数据
def get_daily_data(stock_list, start_date, end_date):
    pro = ts.pro_api('47b4cad5b8693e659ce45d1625d20211311ccb4039d69ba828736eb4')
    all_data = pd.DataFrame()
    for _, stock in stock_list.iterrows():
        ts_code = stock['ts_code']
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        all_data = pd.concat([all_data, df])
        print(f"已获取 {ts_code} 行情数据，共 {len(df)} 条记录")
    return all_data

# 计算技术指标
def calculate_technical_indicators(df):
    result = df.copy()
    result = result.sort_values('trade_date')

    # 1. 简单移动平均线
    result['MA5'] = ta.MA(result['close'], timeperiod=5)
    result['MA10'] = ta.MA(result['close'], timeperiod=10)
    result['MA20'] = ta.MA(result['close'], timeperiod=20)
    result['MA30'] = ta.MA(result['close'], timeperiod=30)
    result['MA60'] = ta.MA(result['close'], timeperiod=60)

    # 2. 指数移动平均线
    result['EMA5'] = ta.EMA(result['close'], timeperiod=5)
    result['EMA10'] = ta.EMA(result['close'], timeperiod=10)
    result['EMA20'] = ta.EMA(result['close'], timeperiod=20)

    # 3. 相对强弱指数
    result['RSI'] = ta.RSI(result['close'], timeperiod=14)

    # 4. MACD
    macd, signal, hist = ta.MACD(result['close'])
    result['MACD'] = macd
    result['MACD_signal'] = signal
    result['MACD_hist'] = hist

    # 5. KDJ指标
    k, d = ta.STOCH(result['high'], result['low'], result['close'])
    result['KDJ_K'] = k
    result['KDJ_D'] = d
    result['KDJ_J'] = 3 * k - 2 * d

    # 6. 布林带
    upper, middle, lower = ta.BBANDS(result['close'], timeperiod=20)
    result['BB_upper'] = upper
    result['BB_middle'] = middle
    result['BB_lower'] = lower
    result['BB_width'] = upper - lower  # 布林带宽度

    # 7. 成交量加权平均价格
    avg_price = ta.AVGPRICE(result['open'], result['high'], result['low'], result['close'])
    result['VWAP'] = (avg_price * result['vol']).cumsum() / result['vol'].cumsum()

    # 8. 平均方向指数
    result['ADX'] = ta.ADX(result['high'], result['low'], result['close'], timeperiod=14)

    # 9. 动量
    result['MOM'] = ta.MOM(result['close'], timeperiod=10)

    # 10. 威廉指标
    result['WILLR'] = ta.WILLR(result['high'], result['low'], result['close'], timeperiod=14)

    # 11. 能量潮
    result['OBV'] = ta.OBV(result['close'], result['vol'])

    # 12. 价格振荡指标
    result['PO'] = PO(result['close'])

    # 13. 慢速随机指标
    slowk, slowd = ta.STOCH(result['high'], result['low'], result['close'])
    result['SlowK'] = slowk
    result['SlowD'] = slowd

    # 14. 快速随机指标
    fastk, fastd = ta.STOCHF(result['high'], result['low'], result['close'])
    result['FastK'] = fastk
    result['FastD'] = fastd

    # 15. 蔡金资金流向
    result['CMF'] = ta.ADOSC(result['high'], result['low'], result['close'], result['vol'])

    # 16. 价格波动率
    result['VOL'] = result['close'].pct_change().rolling(window=20).std() * np.sqrt(20)

    # 17. 相对强弱变化率
    result['RSI_CHG'] = result['RSI'].pct_change()

    # 18. 乖离率
    result['BIAS'] = (result['close'] - result['MA20']) / result['MA20'] * 100

    return result

# 数据打标签
def label_data(df, n=5):
    result = df.copy()
    result = result.sort_values(['ts_code', 'trade_date'])

    result['future_return'] = result.groupby('ts_code')['close'].pct_change(n).shift(-n)

    result['label'] = 0  # 默认标签为0
    result.loc[result['future_return'] > 0.03, 'label'] = 1  # 涨幅超过3%
    result.loc[result['future_return'] < -0.03, 'label'] = -1  # 跌幅超过3%

    # 计算未来5天收益率
    result['return_5d'] = result.groupby('ts_code')['close'].pct_change(5)

    # 计算未来10天收益率
    result['return_10d'] = result.groupby('ts_code')['close'].pct_change(10)

    # 计算未来20天收益率
    result['return_20d'] = result.groupby('ts_code')['close'].pct_change(20)

    return result

# 数据预处理
def preprocess_data(df):
    result = df.copy()
    result = result.dropna(how='all')  # 删除全空的行

    missing_ratio = result.isnull().sum() / len(result)
    columns_to_drop = missing_ratio[missing_ratio > 0.5].index
    result = result.drop(columns_to_drop, axis=1)

    for col in result.columns:
        if result[col].dtype in [np.float64, np.int64]:
            result[col] = result[col].fillna(result[col].median())
        else:
            result[col] = result[col].fillna(result[col].mode()[0])

    for col in result.select_dtypes(include=[np.float64, np.int64]).columns:
        if col not in ['ts_code', 'trade_date', 'label']:
            mean = result[col].mean()
            std = result[col].std()
            lower_bound = mean - 3 * std
            upper_bound = mean + 3 * std
            result.loc[result[col] < lower_bound, col] = lower_bound
            result.loc[result[col] > upper_bound, col] = upper_bound

    return result

# 绘制K线图
def plot_candlestick(df, stock_code, ax=None, show_volume=True, show_patterns=True):
    stock_data = df[df['ts_code'] == stock_code].sort_values('trade_date')

    if ax is None:
        if show_volume:
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10), gridspec_kw={'height_ratios': [3, 1]})
        else:
            fig, ax1 = plt.subplots(1, 1, figsize=(16, 8))
            ax2 = None
        ax = ax1

    # 绘制K线
    for i, row in stock_data.iterrows():
        date = row['trade_date']
        open_price = row['open']
        close_price = row['close']
        high_price = row['high']
        low_price = row['low']

        # 上涨用红色，下跌用绿色
        color = 'red' if close_price >= open_price else 'green'

        # 绘制实体
        ax.bar(date, abs(close_price - open_price), bottom=min(open_price, close_price),
               width=0.6, color=color, alpha=0.8)

        # 绘制上下影线
        ax.vlines(date, low_price, min(open_price, close_price), color=color, linewidth=1)
        ax.vlines(date, max(open_price, close_price), high_price, color=color, linewidth=1)

        # 标记K线形态（如果有）
        if show_patterns:
            patterns = {
                'Morning Star': row['MORNINGSTAR'],
                'Evening Star': row['EVENINGSTAR'],
                'Doji': row['DOJI'],
                'Hammer': row['HAMMER'],
                'Hanging Man': row['HANGINGMAN'],
                'Shooting Star': row['SHOOTINGSTAR'],
                'Dark Cloud': row['DARKCLOUDCOVER'],
                'Bull Engulf': row['ENGULFING_BULL'],
                'Bear Engulf': row['ENGULFING_BEAR']
            }

            for pattern, value in patterns.items():
                if value != 0:
                    if value > 0:
                        # 看涨形态
                        ax.annotate(pattern[:4],
                                    xy=(date, low_price),
                                    xytext=(date, low_price * 0.98),
                                    arrowprops=dict(facecolor='blue', shrink=0.05, width=1.5, headwidth=8),
                                    horizontalalignment='center',
                                    verticalalignment='top',
                                    fontsize=9,
                                    color='blue')
                    else:
                        # 看跌形态
                        ax.annotate(pattern[:4],
                                    xy=(date, high_price),
                                    xytext=(date, high_price * 1.02),
                                    arrowprops=dict(facecolor='red', shrink=0.05, width=1.5, headwidth=8),
                                    horizontalalignment='center',
                                    verticalalignment='bottom',
                                    fontsize=9,
                                    color='red')

    # 绘制移动平均线
    ax.plot(stock_data['trade_date'], stock_data['MA5'], label='MA5', color='blue', linewidth=1.5)
    ax.plot(stock_data['trade_date'], stock_data['MA10'], label='MA10', color='purple', linewidth=1.5)
    ax.plot(stock_data['trade_date'], stock_data['MA20'], label='MA20', color='orange', linewidth=1.5)
    ax.plot(stock_data['trade_date'], stock_data['MA60'], label='MA60', color='green', linewidth=1.5)

    # 添加布林带
    ax.plot(stock_data['trade_date'], stock_data['BB_upper'], label='BB Upper', color='red', linestyle='--',
            linewidth=1)
    ax.plot(stock_data['trade_date'], stock_data['BB_middle'], label='BB Middle', color='gray', linestyle='--',
            linewidth=1)
    ax.plot(stock_data['trade_date'], stock_data['BB_lower'], label='BB Lower', color='red', linestyle='--',
            linewidth=1)
    ax.fill_between(stock_data['trade_date'], stock_data['BB_lower'], stock_data['BB_upper'], color='gray', alpha=0.1)

    ax.set_title(f'{stock_code} Candlestick Chart', fontsize=16)
    ax.set_xlabel('Date', fontsize=12)
    ax.set_ylabel('Price', fontsize=12)
    ax.legend(fontsize=10)
    ax.grid(True, linestyle='--', alpha=0.7)

    # 绘制成交量
    if ax2 is not None:
        for i, row in stock_data.iterrows():
            date = row['trade_date']
            volume = row['vol']
            close_price = row['close']
            open_price = row['open']

            color = 'red' if close_price >= open_price else 'green'
            ax2.bar(date, volume, width=0.6, color=color, alpha=0.8)

        ax2.set_title('Volume', fontsize=14)
        ax2.set_xlabel('Date', fontsize=12)
        ax2.set_ylabel('Volume', fontsize=12)
        ax2.grid(True, linestyle='--', alpha=0.7)

        plt.tight_layout()

    return ax

# 个股画像可视化（雷达图）
def visualize_stock_portrait(df, stock_code, ax=None):
    stock_data = df[df['ts_code'] == stock_code].iloc[-1]  # 取最新数据

    # 选择要展示的特征
    features = [
        'RSI', 'MACD', 'KDJ_K', 'KDJ_D', 'SlowK', 'SlowD',
        'ADX', 'WILLR', 'BIAS', 'VOL', 'BB_width', 'CMF'
    ]

    # 标准化数据
    values = []
    for feature in features:
        if feature in df.columns:
            # 归一化处理
            min_val = df[feature].min()
            max_val = df[feature].max()
            if max_val - min_val != 0:
                value = (stock_data[feature] - min_val) / (max_val - min_val)
            else:
                value = 0.5  # 如果全为相同值，设为0.5
            values.append(value)

    num_features = len(features)
    angles = np.linspace(0, 2 * np.pi, num_features, endpoint=False).tolist()
    angles += angles[:1]
    values += values[:1]

    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 10), subplot_kw={'projection': 'polar'})

    ax.set_theta_zero_location('N')
    ax.set_theta_direction(-1)

    # 绘制雷达图
    ax.plot(angles, values, 'o-', linewidth=2, color='#1f77b4')
    ax.fill(angles, values, alpha=0.25, color='#1f77b4')

    # 添加特征标签
    feature_labels = ['RSI', 'MACD', 'KDJ_K', 'KDJ_D', 'SlowK', 'SlowD',
                     'ADX', 'WILLR', 'BIAS', 'VOL', 'BB_W', 'CMF']
    ax.set_thetagrids([angle * 180 / np.pi for angle in angles[:-1]], feature_labels)

    # 设置刻度范围
    ax.set_ylim(0, 1)

    # 添加标题
    ax.set_title(f'{stock_code} Stock Portrait', fontsize=16)

    # 添加参考线
    for i in range(num_features):
        ax.plot([angles[i], angles[i]], [0, 1], 'k--', alpha=0.1)

    return ax

# 技术指标可视化
def visualize_technical_indicators(df, stock_code):
    stock_data = df[df['ts_code'] == stock_code].sort_values('trade_date')

    fig, axes = plt.subplots(4, 1, figsize=(16, 18), sharex=True)

    # 1. 价格与移动平均线
    axes[0].plot(stock_data['trade_date'], stock_data['close'], label='Close', color='blue', linewidth=2)
    axes[0].plot(stock_data['trade_date'], stock_data['MA5'], label='MA5', color='red', linewidth=1)
    axes[0].plot(stock_data['trade_date'], stock_data['MA10'], label='MA10', color='purple', linewidth=1)
    axes[0].plot(stock_data['trade_date'], stock_data['MA20'], label='MA20', color='orange', linewidth=1)
    axes[0].plot(stock_data['trade_date'], stock_data['MA60'], label='MA60', color='green', linewidth=1)
    axes[0].set_title(f'{stock_code} Price & Moving Averages', fontsize=16)
    axes[0].set_ylabel('Price', fontsize=12)
    axes[0].legend(fontsize=10)
    axes[0].grid(True, linestyle='--', alpha=0.7)

    # 2. MACD指标
    axes[1].plot(stock_data['trade_date'], stock_data['MACD'], label='MACD', color='blue', linewidth=1.5)
    axes[1].plot(stock_data['trade_date'], stock_data['MACD_signal'], label='Signal', color='red', linewidth=1.5)
    axes[1].bar(stock_data['trade_date'], stock_data['MACD_hist'], label='Histogram', color='gray', alpha=0.7)
    axes[1].axhline(y=0, color='black', linestyle='-', alpha=0.3)
    axes[1].set_title('MACD Indicator', fontsize=16)
    axes[1].set_ylabel('MACD Value', fontsize=12)
    axes[1].legend(fontsize=10)
    axes[1].grid(True, linestyle='--', alpha=0.7)

    # 3. KDJ指标
    axes[2].plot(stock_data['trade_date'], stock_data['KDJ_K'], label='KDJ_K', color='blue', linewidth=1.5)
    axes[2].plot(stock_data['trade_date'], stock_data['KDJ_D'], label='KDJ_D', color='red', linewidth=1.5)
    axes[2].plot(stock_data['trade_date'], stock_data['KDJ_J'], label='KDJ_J', color='green', linewidth=1.5)
    axes[2].axhline(y=80, color='gray', linestyle='--', alpha=0.5)
    axes[2].axhline(y=20, color='gray', linestyle='--', alpha=0.5)
    axes[2].set_title('KDJ Indicator', fontsize=16)
    axes[2].set_ylabel('Value', fontsize=12)
    axes[2].legend(fontsize=10)
    axes[2].grid(True, linestyle='--', alpha=0.7)

    # 4. RSI指标
    axes[3].plot(stock_data['trade_date'], stock_data['RSI'], label='RSI', color='purple', linewidth=1.5)
    axes[3].axhline(y=70, color='red', linestyle='--', alpha=0.5)
    axes[3].axhline(y=30, color='green', linestyle='--', alpha=0.5)
    axes[3].set_title('RSI Indicator', fontsize=16)
    axes[3].set_xlabel('Date', fontsize=12)
    axes[3].set_ylabel('RSI Value', fontsize=12)
    axes[3].legend(fontsize=10)
    axes[3].grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.show()

# 多股对比可视化
def compare_stocks(df, stock_codes):
    fig, axes = plt.subplots(2, 1, figsize=(16, 14), sharex=True)

    # 1. 价格走势对比
    for stock_code in stock_codes:
        stock_data = df[df['ts_code'] == stock_code].sort_values('trade_date')
        axes[0].plot(stock_data['trade_date'], stock_data['close'], label=stock_code, linewidth=1.5)

    axes[0].set_title('Stock Price Comparison', fontsize=16)
    axes[0].set_ylabel('Price', fontsize=12)
    axes[0].legend(fontsize=10)
    axes[0].grid(True, linestyle='--', alpha=0.7)

    # 2. 收益率对比
    for stock_code in stock_codes:
        stock_data = df[df['ts_code'] == stock_code].sort_values('trade_date')
        # 计算累积收益率
        stock_data['cum_return'] = (1 + stock_data['close'].pct_change()).cumprod() - 1
        axes[1].plot(stock_data['trade_date'], stock_data['cum_return'], label=stock_code, linewidth=1.5)

    axes[1].set_title('Cumulative Return Comparison', fontsize=16)
    axes[1].set_xlabel('Date', fontsize=12)
    axes[1].set_ylabel('Cumulative Return', fontsize=12)
    axes[1].legend(fontsize=10)
    axes[1].grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.show()

# K线形态分布可视化
def visualize_candlestick_patterns(pattern_counts):
    plt.figure(figsize=(12, 8))

    # 提取形态名称和数量
    patterns = list(pattern_counts.keys())
    counts = list(pattern_counts.values())

    # 创建柱状图
    bars = plt.bar(patterns, counts, color='skyblue')

    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width() / 2., height,
                 f'{height}', ha='center', va='bottom', fontsize=12)

    # 设置图表标题和标签
    plt.title('Candlestick Pattern Distribution', fontsize=16)
    plt.xlabel('Pattern', fontsize=12)
    plt.ylabel('Count', fontsize=12)
    plt.xticks(rotation=45, ha='right', fontsize=10)
    plt.grid(axis='y', linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.show()

# 数据可视化主函数
def visualize_data(df, stock_code=None):
    if stock_code:
        # 个股K线图
        plot_candlestick(df, stock_code)
        plt.show()

        # 技术指标可视化
        visualize_technical_indicators(df, stock_code)

        # 个股画像雷达图
        visualize_stock_portrait(df, stock_code)
        plt.show()

        # K线形态识别和可视化
        stock_data, pattern_counts = identify_candlestick_patterns(df[df['ts_code'] == stock_code])
        visualize_candlestick_patterns(pattern_counts)

        # 标签分布
        plt.figure(figsize=(10, 6))
        ax = sns.countplot(x='label', data=df)
        plt.title('Label Distribution', fontsize=16)
        plt.xlabel('Label', fontsize=12)
        plt.ylabel('Count', fontsize=12)
        plt.xticks(rotation=45)

        # 添加数据标签
        for p in ax.patches:
            ax.annotate(f'{p.get_height()}', (p.get_x() + p.get_width() / 2., p.get_height()),
                        ha='center', va='center', xytext=(0, 10), textcoords='offset points')
        plt.tight_layout()
        plt.show()

    # 多股对比（选择前5只股票）
    if len(df['ts_code'].unique()) > 1:
        compare_stocks(df, df['ts_code'].unique()[:5])

    # 相关性分析
    numeric_cols = df.select_dtypes(include=[np.float64, np.int64]).columns
    corr = df[numeric_cols].corr()

    plt.figure(figsize=(18, 16))
    mask = np.triu(np.ones_like(corr, dtype=bool))
    cmap = sns.diverging_palette(230, 20, as_cmap=True)
    sns.heatmap(corr, mask=mask, annot=True, cmap=cmap, fmt='.2f',
                annot_kws={'size': 8}, square=True, linewidths=.5, cbar_kws={"shrink": .8})
    plt.title('Feature Correlation Heatmap', fontsize=16)
    plt.tight_layout()
    plt.show()

    # 主成分分析
    pca = PCA(n_components=2)
    X = df.select_dtypes(include=[np.float64, np.int64])
    X = X.drop(['label'], axis=1, errors='ignore')

    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    X_pca = pca.fit_transform(X_scaled)

    plt.figure(figsize=(14, 12))
    scatter = plt.scatter(X_pca[:, 0], X_pca[:, 1], c=df['label'],
                          cmap='viridis', alpha=0.7, s=60, edgecolor='k', linewidth=0.5)
    plt.title('PCA Result', fontsize=16)
    plt.xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.2%})', fontsize=12)
    plt.ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.2%})', fontsize=12)

    # 添加颜色条和标签
    cbar = plt.colorbar(scatter)
    cbar.set_label('Label', fontsize=12)
    cbar.set_ticks([-1, 0, 1])
    cbar.set_ticklabels(['Down', 'Neutral', 'Up'])

    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()

# 构建机器学习模型
def build_model(df):
    X = df.select_dtypes(include=[np.float64, np.int64])
    X = X.drop(['label', 'future_return'], axis=1, errors='ignore')
    y = df['label']

    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.3, random_state=42)

    model = RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=-1)
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)

    accuracy = accuracy_score(y_test, y_pred)
    report = classification_report(y_test, y_pred)
    cm = confusion_matrix(y_test, y_pred)

    # 可视化混淆矩阵
    plt.figure(figsize=(14, 12))
    ax = sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                     xticklabels=['Down', 'Neutral', 'Up'],
                     yticklabels=['Down', 'Neutral', 'Up'],
                     annot_kws={'size': 14})
    plt.title('Confusion Matrix', fontsize=16)
    plt.xlabel('Predicted Label', fontsize=14)
    plt.ylabel('True Label', fontsize=14)

    # 添加标签和百分比
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            text = ax.texts[i * cm.shape[1] + j]
            text.set_text(f"{text.get_text()}\n({cm[i, j] / np.sum(cm[i]) * 100:.1f}%)")

    plt.tight_layout()
    plt.show()

    # 可视化特征重要性
    feature_importance = model.feature_importances_
    feature_names = X.columns

    plt.figure(figsize=(16, 14))
    indices = np.argsort(feature_importance)[-20:]  # 取前20个重要特征
    plt.barh(range(len(indices)), feature_importance[indices], align='center', color='#1f77b4')
    plt.yticks(range(len(indices)), [feature_names[i] for i in indices], fontsize=12)
    plt.xlabel('Feature Importance', fontsize=14)
    plt.title('Random Forest Feature Importance', fontsize=16)
    plt.grid(axis='x', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()

    return {
        'model': model,
        'accuracy': accuracy,
        'report': report,
        'feature_importance': feature_importance,
        'feature_names': feature_names
    }

# 主函数
def main():
    end_date = datetime.now().strftime('%Y%m%d')
    start_date = (datetime.now() - timedelta(days=730)).strftime('%Y%m%d')

    print("正在获取股票列表...")
    stock_list = get_stock_list()
    selected_stocks = stock_list.head(10)
    print(f"已获取 {len(selected_stocks)} 支股票")

    print("正在获取行情数据...")
    daily_data = get_daily_data(selected_stocks, start_date, end_date)
    print(f"共获取 {len(daily_data)} 条行情数据")

    print("正在计算技术指标...")
    technical_data = calculate_technical_indicators(daily_data)

    print("正在识别K线形态...")
    technical_data, _ = identify_candlestick_patterns(technical_data)

    print("正在为数据打标签...")
    labeled_data = label_data(technical_data)

    print("正在进行数据预处理...")
    preprocessed_data = preprocess_data(labeled_data)

    print("正在进行数据可视化...")
    sample_stock = selected_stocks.iloc[0]['ts_code']
    visualize_data(preprocessed_data, sample_stock)

    print("正在构建机器学习模型...")
    model_results = build_model(preprocessed_data)

    print(f"模型准确率: {model_results['accuracy']:.4f}")
    print("分类报告:")
    print(model_results['report'])

if __name__ == "__main__":
    main()